import os
import re
import tkinter as tk
from tkinter import filedialog, messagebox
from pathlib import Path

class MediaFileProcessor:
    def __init__(self):
        self.obsidian_root = None
        self.media_dir = None
        self.output_dir = None
        
    def select_directories(self):
        """弹窗选择三个目录"""
        root = tk.Tk()
        root.withdraw()  # 隐藏主窗口
        
        # 选择Obsidian库根目录
        self.obsidian_root = filedialog.askdirectory(
            title="选择Obsidian库的根目录"
        )
        if not self.obsidian_root:
            return False
        
        # 选择媒体文件目录（必须是Obsidian库的子目录）
        self.media_dir = filedialog.askdirectory(
            title="选择媒体文件所在的目录（将作为tags来源）",
            initialdir=self.obsidian_root
        )
        if not self.media_dir:
            return False
        
        # 验证媒体目录是否是Obsidian库的子目录
        try:
            if not Path(self.media_dir).is_relative_to(Path(self.obsidian_root)):
                messagebox.showerror("错误", "媒体文件目录必须是Obsidian库的子目录！")
                return False
        except AttributeError:
            # 对于Python 3.8以下的版本，使用替代方法
            if not self.media_dir.startswith(self.obsidian_root):
                messagebox.showerror("错误", "媒体文件目录必须是Obsidian库的子目录！")
                return False
        
        # 选择md文件输出目录（必须是Obsidian库的子目录）
        self.output_dir = filedialog.askdirectory(
            title="选择md文件输出目录",
            initialdir=self.obsidian_root
        )
        if not self.output_dir:
            return False
        
        # 验证输出目录是否是Obsidian库的子目录
        try:
            if not Path(self.output_dir).is_relative_to(Path(self.obsidian_root)):
                messagebox.showerror("错误", "输出目录必须是Obsidian库的子目录！")
                return False
        except AttributeError:
            if not self.output_dir.startswith(self.obsidian_root):
                messagebox.showerror("错误", "输出目录必须是Obsidian库的子目录！")
                return False
        
        return True
    
    def scan_all_files(self):
        """扫描目录下所有文件"""
        if not self.media_dir:
            return []
        
        all_files = []
        for root, dirs, files in os.walk(self.media_dir):
            for file in files:
                file_path = Path(root) / file
                all_files.append(file_path)
        
        return all_files
    
    def parse_filename(self, filename):
        """
        解析文件名，提取番号、女优名、影片名
        支持两种格式：
        1. 番号+女优名+影片名+后缀，如：ABP-418 彩美旬果《社内恋愛SEX》2.jpg
        2. 番号+影片名+女优名+后缀
        
        如果无法解析影片名，直接使用原文件名
        """
        # 移除扩展名
        name_only = filename.stem
        
        # 匹配影片名（书名号内的内容）
        title_match = re.search(r'《([^》]+)》', name_only)
        if title_match:
            title = title_match.group(1)
        else:
            # 如果无法解析影片名，直接使用原文件名（不含扩展名）
            title = name_only
        
        # 尝试提取番号（通常以字母开头，包含数字和短横线）
        id_match = re.search(r'^([A-Za-z]+-\d+|[A-Za-z]+\d+)', name_only.strip())
        if id_match:
            media_id = id_match.group(1)
        else:
            # 如果没有番号，使用文件名前20个字符作为ID（避免过长）
            media_id = name_only[:20].replace(' ', '_')
        
        # 尝试提取演员名
        actor = "未知演员"
        if title_match:
            # 移除影片名部分，便于提取其他信息
            name_without_title = re.sub(r'《[^》]+》', '', name_only)
            
            if id_match:
                # 移除番号部分
                actor_part = re.sub(r'^' + re.escape(media_id), '', name_without_title).strip()
                # 清理数字和其他特殊字符
                actor_candidate = re.sub(r'\s*\d+$', '', actor_part).strip()
                if actor_candidate:
                    actor = actor_candidate
            else:
                # 如果没有番号，尝试从文件名中提取可能的演员名
                # 假设演员名是中文字符串
                chinese_words = re.findall(r'[\u4e00-\u9fff]{2,}', name_without_title)
                if chinese_words:
                    actor = chinese_words[0]
        else:
            # 如果没有书名号，尝试从整个文件名中提取可能的演员名
            chinese_words = re.findall(r'[\u4e00-\u9fff]{2,}', name_only)
            if chinese_words:
                actor = chinese_words[0]
        
        print(f"文件解析: {filename.name}")
        print(f"  番号: {media_id}, 影片名: {title}, 女优: {actor}")
        
        return {
            'id': media_id,
            'title': title,
            'actor': actor,
            'filename': filename.name,
            'file_path': filename,
            'parsed_successfully': bool(title_match and id_match)  # 标记是否成功解析
        }
    
    def get_obsidian_path(self, file_path):
        """获取Obsidian格式的路径（使用正斜杠）"""
        try:
            relative_path = Path(file_path).relative_to(self.obsidian_root)
            # 将Windows路径分隔符转换为Obsidian使用的正斜杠
            obsidian_path = str(relative_path).replace('\\', '/')
            return obsidian_path
        except ValueError:
            # 如果路径不在Obsidian根目录下，返回文件名
            return Path(file_path).name
    
    def get_tags(self, media_info):
        """生成tags：包含actors和所选目录名（不包含父目录名）"""
        tags = []
        
        # 添加演员作为tag
        if media_info['actor'] and media_info['actor'] != "未知演员":
            tags.append(f"#{media_info['actor']}")
        
        # 添加所选目录名作为tag（只取最后一级目录名）
        if self.media_dir:
            dir_name = Path(self.media_dir).name
            if dir_name and dir_name != '.':
                tags.append(f"#{dir_name}")
                
                # 目录名称拆词
                words = re.findall(r'[\w\u4e00-\u9fff]+', dir_name)
                for word in words:
                    if len(word) > 1 and word not in media_info.get('actor', ''):
                        tags.append(f"#{word}")
        
        return list(set(tags))  # 去重
    
    def yaml_dump(self, data):
        """简单的YAML格式输出函数（替代pyyaml）"""
        lines = []
        for key, value in data.items():
            if isinstance(value, list):
                if value:
                    lines.append(f"{key}:")
                    for item in value:
                        lines.append(f"  - {item}")
                else:
                    lines.append(f"{key}: []")
            else:
                lines.append(f"{key}: {value}")
        return "\n".join(lines)
    
    def parse_front_matter(self, content):
        """解析front matter内容"""
        if not content.startswith('---'):
            return None, content
        
        parts = content.split('---', 2)
        if len(parts) < 3:
            return None, content
        
        front_matter_str = parts[1]
        body_content = parts[2] if len(parts) > 2 else ""
        
        front_matter = {}
        lines = front_matter_str.strip().split('\n')
        current_key = None
        
        for line in lines:
            line = line.strip()
            if not line:
                continue
                
            if line.startswith('- '):
                # 列表项
                if current_key and current_key in front_matter:
                    if isinstance(front_matter[current_key], list):
                        front_matter[current_key].append(line[2:].strip())
            elif ':' in line:
                # 新键值对
                key, value = line.split(':', 1)
                key = key.strip()
                value = value.strip()
                
                if value == '':
                    # 空值，可能是列表的开始
                    front_matter[key] = []
                    current_key = key
                else:
                    front_matter[key] = value
                    current_key = None
            elif current_key and current_key in front_matter and isinstance(front_matter[current_key], list):
                # 续行（简化处理）
                pass
        
        return front_matter, body_content.strip()
    
    def get_existing_md_content(self, md_file_path):
        """读取现有的md文件内容，提取需要保留的部分"""
        if not md_file_path.exists():
            return None, ""
        
        try:
            with open(md_file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            return self.parse_front_matter(content)
        except Exception as e:
            print(f"读取现有md文件失败: {e}")
            return None, content
    
    def create_md_content(self, media_info, media_files, existing_front_matter=None, existing_body=""):
        """创建或更新Markdown文件内容"""
        # 基础front matter
        tags = self.get_tags(media_info)
        front_matter = {
            'tags': tags,
            'type': 'movie',
            'id': media_info['id'],
            'title': media_info['title'],
            'actors': []
        }
        
        # 合并现有的actors信息
        if existing_front_matter and 'actors' in existing_front_matter:
            if isinstance(existing_front_matter['actors'], list):
                front_matter['actors'] = existing_front_matter['actors']
            else:
                front_matter['actors'] = [existing_front_matter['actors']]
        
        # 添加当前演员（如果不存在且不是"未知演员"）
        if (media_info['actor'] and media_info['actor'] != "未知演员" and 
            media_info['actor'] not in front_matter['actors']):
            front_matter['actors'].append(media_info['actor'])
        
        # 合并现有的tags信息
        if existing_front_matter and 'tags' in existing_front_matter:
            if isinstance(existing_front_matter['tags'], list):
                existing_tags = existing_front_matter['tags']
            else:
                existing_tags = [existing_front_matter['tags']]
            
            # 去重合并
            all_tags = list(set(front_matter['tags'] + existing_tags))
            front_matter['tags'] = all_tags
        
        # 按文件扩展名分类
        image_files = []
        video_files = []
        other_files = []
        
        image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp'}
        video_extensions = {'.mp4', '.avi', '.mov', '.wmv', '.flv', '.mkv', '.webm'}
        
        for file in media_files:
            ext = file.suffix.lower()
            if ext in image_extensions:
                image_files.append(file)
            elif ext in video_extensions:
                video_files.append(file)
            else:
                other_files.append(file)
        
        # 构建新的正文内容
        body_content = ""
        
        # 图片部分
        if image_files:
            body_content += "# 图片\n"
            for img_file in sorted(image_files):
                obsidian_path = self.get_obsidian_path(img_file)
                body_content += f"![[{obsidian_path}]]\n"
            body_content += "\n"
        
        # 视频部分
        if video_files:
            body_content += "# 视频\n"
            for vid_file in sorted(video_files):
                obsidian_path = self.get_obsidian_path(vid_file)
                body_content += f"![[{obsidian_path}]]\n"
            body_content += "\n"
        
        # 其他文件部分
        if other_files:
            body_content += "# 其他文件\n"
            for other_file in sorted(other_files):
                obsidian_path = self.get_obsidian_path(other_file)
                body_content += f"![[{obsidian_path}]]\n"
            body_content += "\n"
        
        # 如果现有正文中有图片或视频部分，用新的替换
        if existing_body:
            # 移除现有的图片、视频和其他文件部分
            lines = existing_body.split('\n')
            new_lines = []
            in_media_section = False
            
            for line in lines:
                if line.strip() in ["# 图片", "# 视频", "# 其他文件"]:
                    in_media_section = True
                    continue
                elif line.strip().startswith("# ") and in_media_section:
                    in_media_section = False
                
                if not in_media_section:
                    if line.strip() and not line.strip().startswith("!["):
                        new_lines.append(line)
            
            # 添加保留的内容
            if new_lines:
                body_content = "\n".join(new_lines) + "\n\n" + body_content
        
        # 组合完整内容
        content = "---\n"
        content += self.yaml_dump(front_matter)
        content += "\n---\n\n"
        content += body_content
        
        return content
    
    def process_files(self):
        """主处理函数"""
        if not self.select_directories():
            print("目录选择取消或出错，程序退出")
            return
        
        print(f"Obsidian根目录: {self.obsidian_root}")
        print(f"媒体文件目录（tags来源）: {self.media_dir}")
        print(f"输出目录: {self.output_dir}")
        
        print("开始扫描所有文件...")
        all_files = self.scan_all_files()
        print(f"找到 {len(all_files)} 个文件")
        
        if not all_files:
            messagebox.showinfo("提示", "未找到任何文件！")
            return
        
        # 按番号分组
        media_groups = {}
        unparsed_files = []
        
        for file_path in all_files:
            media_info = self.parse_filename(file_path)
            if media_info:
                media_id = media_info['id']
                if media_id not in media_groups:
                    media_groups[media_id] = {
                        'info': media_info,
                        'files': []
                    }
                media_groups[media_id]['files'].append(file_path)
                
                if not media_info['parsed_successfully']:
                    unparsed_files.append(file_path)
        
        print(f"识别出 {len(media_groups)} 个不同的番号/文件组")
        
        if unparsed_files:
            print(f"\n以下 {len(unparsed_files)} 个文件未能完全解析:")
            for file_path in unparsed_files:
                print(f"  {file_path.name}")
        
        # 创建或更新Markdown文件
        created_count = 0
        updated_count = 0
        
        for media_id, group in media_groups.items():
            md_file_path = Path(self.output_dir) / f"{media_id}.md"
            
            existing_front_matter = None
            existing_body = ""
            
            if md_file_path.exists():
                existing_front_matter, existing_body = self.get_existing_md_content(md_file_path)
                updated_count += 1
                action = "更新"
            else:
                created_count += 1
                action = "创建"
            
            content = self.create_md_content(
                group['info'], 
                group['files'], 
                existing_front_matter, 
                existing_body
            )
            
            try:
                with open(md_file_path, 'w', encoding='utf-8') as f:
                    f.write(content)
                print(f"{action}文件: {md_file_path.name}")
                
                # 显示文件统计
                image_count = len([f for f in group['files'] if f.suffix.lower() in {'.jpg', '.jpeg', '.png', '.gif', '.bmp'}])
                video_count = len([f for f in group['files'] if f.suffix.lower() in {'.mp4', '.avi', '.mov', '.wmv', '.flv', '.mkv', '.webm'}])
                other_count = len(group['files']) - image_count - video_count
                print(f"  包含: {image_count}图片, {video_count}视频, {other_count}其他文件")
                
            except Exception as e:
                print(f"写入文件失败 {md_file_path.name}: {e}")
        
        # 显示结果摘要
        summary = (f"处理完成！\n"
                  f"扫描文件总数: {len(all_files)}\n"
                  f"创建: {created_count} 个md文件\n"
                  f"更新: {updated_count} 个md文件\n"
                  f"未完全解析文件: {len(unparsed_files)} 个")
        
        print(f"\n{summary}")
        messagebox.showinfo("处理完成", summary)

def main():
    processor = MediaFileProcessor()
    processor.process_files()

if __name__ == "__main__":
    main()